import sympy as sy
import numpy as np
import matplotlib.pyplot as plt

def f_drag(v_vec, m, g, v_terminal):
    """
    Calcula a força de arrasto (resistência do ar) sobre a bola.

    Parâmetros:
        v_vec (np.ndarray): vetor de velocidade (3D).
        m (float): massa da bola (kg).
        g (float): aceleração da gravidade (m/s²).
        v_terminal (float): velocidade terminal (m/s).

    Retorna:
        np.ndarray: vetor força de arrasto (3D).
    """
    return -(m * g / v_terminal**2) * np.linalg.norm(v_vec) * v_vec


def f_magnus(v_vec, omega_vec, r, A, rho):
    """
    Calcula a força de Magnus devido à rotação da bola.

    Parâmetros:
        v_vec (np.ndarray): vetor de velocidade (3D).
        omega_vec (np.ndarray): vetor de rotação (rad/s).
        rho (float): densidade do ar (kg/m³).
        r (float): raio da bola (m).
        A (float): área frontal da bola (m²).

    Retorna:
        np.ndarray: vetor força de Magnus (3D).
    """
    return 0.5 * A * rho * r * np.cross(omega_vec, v_vec)

def simular_trajetoria_3d(posicao_inicial, velocidade_inicial, omega_vec, m, r, A, vT, dt=0.001, tempo_total=5.0):
    """
    Simula o movimento 3D da bola com o método de Euler, incluindo forças de gravidade,
    arrasto e Magnus.

    Parâmetros:
        posicao_inicial (np.ndarray): posição inicial da bola (3D).
        velocidade_inicial (np.ndarray): velocidade inicial da bola (3D).
        omega_vec (np.ndarray): vetor de rotação (rad/s).
        m (float): massa da bola (kg).
        r (float): raio da bola (m).
        A (float): área frontal da bola (m²).
        vT (float): velocidade terminal (m/s).
        dt (float): passo de tempo (s).
        tempo_total (float): duração total da simulação (s).

    Retorna:
        np.ndarray: trajetória da bola como array de posições.
    """
    g = 9.81  # m/s²
    r_vec = posicao_inicial.copy()
    v_vec = velocidade_inicial.copy()
    trajetoria = [r_vec.copy()]
    n_passos = int(tempo_total / dt)

    for _ in range(n_passos):
        F_drag = f_drag(v_vec, m, g, vT)
        F_magnus = f_magnus(v_vec, omega_vec, r, A)
        F_grav = np.array([0.0, -m * g, 0.0])
        F_total = F_drag + F_magnus + F_grav
        a_vec = F_total / m

        v_vec += a_vec * dt
        r_vec += v_vec * dt

        if r_vec[2] < -5:  # Limite inferior do campo
            break
        trajetoria.append(r_vec.copy())

    return np.array(trajetoria)





def plotar_trajetoria_3d(trajetoria):
    """
    Plota a trajetória 3D da bola.

    Parâmetros:
        trajetoria (np.ndarray): array com posições 3D da bola ao longo do tempo.
    """
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(trajetoria[:, 0], trajetoria[:, 1], trajetoria[:, 2], label='Trajetória da bola')
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Z (m)')
    ax.set_title('Movimento 3D da Bola com Efeito Magnus')
    ax.legend()
    plt.show()





def verificar_golo(trajetoria):
    """
    Verifica se a bola entrou na baliza.

    Condições da baliza:
        x <= 0
        0 < y < 2.4
        -3.66 < z < 3.66

    Parâmetros:
        trajetoria (np.ndarray): array com posições 3D da bola ao longo do tempo.

    Retorna:
        tuple: (ponto mais próximo, booleano indicando se foi golo)
    """
    x_coords = trajetoria[:, 0]
    idx_closest_to_x0 = np.argmin(np.abs(x_coords))
    ponto = trajetoria[idx_closest_to_x0]
    x, y, z = ponto
    entrou = (x <= 0) and (0 < y < 2.4) and (-3.66 < z < 3.66)
    return ponto, entrou






# ======================
# EXEMPLO DE USO
# ======================

if __name__ == "__main__":
    # Definições iniciais
    m = 0.45
    r = 0.11
    A = np.pi * r**2
    vT = 100 / 3.6  # 100 km/h para m/s
    omega_y = 390
    omega_vec = np.array([0.0, omega_y, 0.0])
    r_inicial = np.array([0.0, 0.0, 23.8])
    v_inicial = np.array([25.0, 5.0, -50.0])

    # Simulação
    traj = simular_trajetoria_3d(r_inicial, v_inicial, omega_vec, m, r, A, vT)

    # Visualização
    plotar_trajetoria_3d(traj)

    # Verificação de golo
    ponto_golo, foi_golo = verificar_golo(traj)
    print("Ponto mais próximo do gol:", ponto_golo)
    print("Entrou na baliza?", foi_golo)